package org.example;

import com.google.common.collect.Ordering;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;

import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.stream.Collectors;

@Slf4j
public class Heap implements ISwap {

    @Test
    public synchronized void heap() {
        /* 初始化堆 */
        // 初始化小顶堆
        Queue<Integer> minHeap = new PriorityQueue<>(3);
        /* 输入列表并建堆 */
        minHeap = new PriorityQueue<>(Arrays.asList(1, 3, 2, 5, 4));
        // 初始化大顶堆（使用 lambda 表达式修改 Comparator 即可）
        Queue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);

        /* 元素入堆 */
        maxHeap.offer(1);
        maxHeap.offer(3);
        maxHeap.offer(2);
        maxHeap.offer(5);
        maxHeap.offer(4);
        maxHeap.offer(4);
        maxHeap.offer(4);

        while (!maxHeap.isEmpty()) {
            log.info(maxHeap.poll() + "");
        }
    }

    @Test
    void topK() {

        @RequiredArgsConstructor
        class TopK<E extends Comparable>
                extends PriorityQueue<E> {

            int maxSize;

            public TopK(int maxSize) {
                super((o1, o2) -> Ordering.natural().compare(o1, o2));
                this.maxSize = maxSize;
            }

            @Override
            public boolean add(E e) {
                if (this.size() >= maxSize) {
                    E peek = peek();
                    int compare = e.compareTo(peek);
                    if (compare > 0) {
                        poll();
                    } else {
                        return false;
                    }
                }

                return super.add(e);
            }
        }

        TopK<Integer> maxHeap = new TopK<>(2);

        /* 元素入堆 */
        maxHeap.add(1);
        maxHeap.add(3);
        maxHeap.add(2);
        maxHeap.add(5);
        maxHeap.add(4);

        List<Integer> collect = maxHeap.stream().sorted((o1, o2) -> o2 - o1).collect(Collectors.toList());
        for (Integer i : collect) {
            log.info(i + "");
        }
        while (!maxHeap.isEmpty()) {
            log.info(maxHeap.poll() + "");
        }

    }
}
